import numpy as np

import manim
from .pi_creature import Yan, PiCreatureGroup
from .svg_drawing import SpeechBubble

__all__ = [
    'Blink',
    'Emo',
    'LookAt',
    "PiCreatureBubbleIntroduction",
    "PiCreatureSays",
    "RemovePiCreatureBubble",
    'FlashThroughClass'
]


class Blink(manim.ApplyMethod):
    def __init__(self, pi_creature, rate_func=manim.squish_rate_func(manim.there_and_back), **kwargs):
        manim.ApplyMethod.__init__(self, pi_creature.blink, rate_func=rate_func, **kwargs)


class Emo(manim.ApplyMethod):
    def __init__(self, pi_creature, happy_or_sad: float, run_time=0.5, **kwargs):
        """
        :param pi_creature:
        :param happy_or_sad: happy:0.1~+1 sad:0.1~-1 peace:0
        :param kwargs: rate_func=linear, run_time=0.5
        """
        # ApplyMethod.__init__(self, pi_creature.smile, False, rate_func=linear, run_time=0.5, **kwargs)
        manim.ApplyMethod.__init__(self, pi_creature.emo, happy_or_sad,
                                   rate_func=manim.linear, run_time=run_time, **kwargs)


class LookAt(manim.ApplyMethod):
    def __init__(self, pi_creature, thing_to_look_at, rate_func=manim.linear, run_time=0.5, **kwargs):
        assert isinstance(pi_creature, Yan)
        super().__init__(pi_creature.look_at, thing_to_look_at, rate_func=rate_func, run_time=run_time, **kwargs)


def creature_bubble_animations(my_pi, *content, bubble_class=SpeechBubble, **kwargs):
    added_anims = kwargs.pop("added_anims", [])
    bubble_kwargs = kwargs.pop("bubble_kwargs", {})
    bubble_removal_kwargs = kwargs.pop("bubble_removal_kwargs", {})
    anims = []

    def has_bubble(pi):
        return hasattr(pi, "bubble") and pi.bubble is not None

    if has_bubble(my_pi):
        old_bubble = my_pi.bubble
        bubble = my_pi.get_bubble(
            *content,
            bubble_class=bubble_class
        )
        anims += [
            manim.ReplacementTransform(old_bubble, bubble),
            manim.ReplacementTransform(old_bubble.content, bubble.content),
            # pi_creature.change_mode, target_mode
        ]
    else:
        anims.append(PiCreatureBubbleIntroduction(
            my_pi,
            *content,
            bubble_class=bubble_class,
            bubble_kwargs=bubble_kwargs,
            **kwargs
        ))
    anims += added_anims
    return anims


class PiCreatureBubbleIntroduction(manim.AnimationGroup):
    def __init__(self, pi_creature, *content, target_mode="speaking", bubble_class=SpeechBubble, change_mode_kwargs={},
                 bubble_creation_class=manim.DrawBorderThenFill, bubble_creation_kwargs={}, bubble_kwargs={},
                 content_introduction_class=manim.Write, content_introduction_kwargs={}, look_at_arg=None, **kwargs):
        self.target_mode = target_mode
        self.bubble_class = bubble_class
        self.change_mode_kwargs = change_mode_kwargs
        self.bubble_creation_class = bubble_creation_class
        self.bubble_creation_kwargs = bubble_creation_kwargs
        self.bubble_kwargs = bubble_kwargs
        self.content_introduction_class = content_introduction_class
        self.content_introduction_kwargs = content_introduction_kwargs
        self.look_at_arg = look_at_arg
        bubble = pi_creature.get_bubble(
            *content,
            bubble_class=self.bubble_class,
            **self.bubble_kwargs
        )
        manim.Group(bubble, bubble.content).shift_onto_screen()

        # TODO: no change mode
        # pi_creature.generate_target()
        # pi_creature.target.change_mode(self.target_mode)
        # if self.look_at_arg is not None:
        #     pi_creature.target.look_at(self.look_at_arg)

        # change_mode = manim.MoveToTarget(pi_creature, **self.change_mode_kwargs)
        bubble_creation = self.bubble_creation_class(
            bubble, **self.bubble_creation_kwargs
        )
        content_introduction = self.content_introduction_class(
            bubble.content, **self.content_introduction_kwargs
        )
        # FIXME remove change_mode
        manim.AnimationGroup.__init__(
            self, bubble_creation, content_introduction,
            **kwargs
        )


class PiCreatureSays(PiCreatureBubbleIntroduction):
    def __init__(self, pi_creator, *content, target_mode="speaking", bubble_class: object = SpeechBubble, **kwargs):
        super().__init__(pi_creator, *content, target_mode=target_mode, bubble_class=bubble_class, **kwargs)


class RemovePiCreatureBubble(manim.AnimationGroup):
    def __init__(self, pi_creature, target_mode="plain", look_at_arg=None, remover=True, **kwargs):
        assert hasattr(pi_creature, "bubble")
        self.target_mode = target_mode
        self.look_at_arg = look_at_arg
        self.remover = remover
        self.pi_creature = pi_creature
        # pi_creature.generate_target()
        # pi_creature.target.change_mode(self.target_mode)
        # if self.look_at_arg is not None:
        #     pi_creature.target.look_at(self.look_at_arg)

        manim.AnimationGroup.__init__(
            self,
            # manim.MoveToTarget(pi_creature),
            manim.FadeOut(pi_creature.bubble),
            manim.FadeOut(pi_creature.bubble.content),
        )

    def clean_up_from_scene(self, scene=None):
        manim.AnimationGroup.clean_up_from_scene(self, scene)
        self.pi_creature.bubble = None
        if scene is not None:
            scene.add(self.pi_creature)


class FlashThroughClass(manim.Animation):
    def __init__(self, mobject, mode="linear", highlight_color=manim.GREEN, **kwargs):
        self.highlight_color = highlight_color
        if not isinstance(mobject, PiCreatureGroup):
            raise Exception("FlashThroughClass mobject must be a PiCreatureClass")
        self.indices = list(range(mobject.col_height * mobject.row_width))
        if mode == "random":
            np.random.shuffle(self.indices)
        manim.Animation.__init__(self, mobject, **kwargs)

    def interpolate_mobject(self, alpha):
        index = int(np.floor(alpha * self.mobject.col_height * self.mobject.row_width))
        for pi in self.mobject:
            pi.set_color(manim.BLUE_E)
        if index < self.mobject.col_height * self.mobject.row_width:
            self.mobject[self.indices[index]].set_color(self.highlight_color)
